package net.csiit.platform.codeguard.maven.plugin;

import net.csiit.platform.codeguard.tools.AES;
import net.csiit.platform.codeguard.tools.Constants;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.maven.plugin.AbstractMojo;
import org.apache.maven.plugins.annotations.Parameter;

import java.io.*;
import java.util.Enumeration;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;
import java.util.zip.ZipEntry;

/**
 * @author adam
 * @date 2020/7/1
 */
public abstract class AbstractGuardPackageMojo extends AbstractMojo {

    public static final String GUARD_WORKSPACE = "guard";
    public static final String CLASS_TEMP_DIR = "classes";
    public static final String DEFAULT_POSTFIX = GUARD_WORKSPACE;

    /**
     * target directory
     */
    @Parameter(defaultValue = "${project.build.directory}")
    protected File target;

    protected File guardWorkspace;

    protected File classTempDir;

    /**
     * timestamp format
     */
    @Parameter(defaultValue = "${maven.build.timestamp.format}")
    private String format;

    /**
     * get jar file
     *
     * @param directory:   target directory
     * @param jarFileName: specify the jar file name. If null, then return first jar file
     * @return
     */
    protected JarFile getJarFile(File directory, String jarFileName) throws IOException {
        if (directory == null || !directory.isDirectory()) {
            return null;
        }

        File[] files = directory.listFiles(new FilenameFilter() {
            @Override
            public boolean accept(File dir, String name) {
                if (StringUtils.isEmpty(jarFileName)) {
                    return StringUtils.endsWith(name, ".jar");
                }

                return StringUtils.equals(name, jarFileName) || name.matches(jarFileName);
            }
        });

        if (ArrayUtils.isEmpty(files)) {
            return null;
        }

        return new JarFile(files[0], true);
    }


    /**
     * get or create code guard workspace directory
     *
     * @return
     */
    protected synchronized File getGuardWorkspace() {
        if (guardWorkspace == null) {
            guardWorkspace = new File(target.getAbsolutePath(), GUARD_WORKSPACE);
            if (!guardWorkspace.exists()) {
                guardWorkspace.mkdir();
            }
        }

        return guardWorkspace;
    }

    /**
     * class file temporary directory
     *
     * @return
     */
    protected synchronized File getClassTempDir() {
        if (classTempDir == null) {
            File workspace = getGuardWorkspace();
            classTempDir = new File(workspace, CLASS_TEMP_DIR);
            if (!classTempDir.exists()) {
                classTempDir.mkdir();
            }
        }

        return classTempDir;
    }

    /**
     * encrypt a jar file using specified key, and rename source jar file.
     *
     * @param sourceJar the jar file needed to be encrypt
     * @param key       the encrypt key
     * @param postfix   postfix to rename source jar file
     * @param encrypt   true for encrypt jar, false to decrypt jar
     * @return the encrypted jar file
     */
    protected File handleJarFile(JarFile sourceJar, String key, String postfix, List<String> handlePackages, boolean encrypt) throws IOException {
        if (sourceJar == null) {
            return null;
        }

        if (StringUtils.isEmpty(key)) {
            return null;
        }

        if (StringUtils.isEmpty(postfix)) {
            postfix = DEFAULT_POSTFIX;
        }

        Manifest manifest = sourceJar.getManifest();
        Enumeration<JarEntry> entryEnumeration = sourceJar.entries();
        String fileName = sourceJar.getName().substring(sourceJar.getName().lastIndexOf(File.separator) + 1);

        Manifest targetFest = new Manifest(manifest);
        targetFest.getMainAttributes().putValue(Constants.GUARD_MANIFEST_ATTR, Constants.GUARD_MANIFEST_VALUE);

        File targetFile = createTempTargetFile(target.getAbsolutePath(), fileName.replace(".jar", "").concat(".").concat(postfix).concat(".jar"));
        try (OutputStream outputStream = new FileOutputStream(targetFile);
             JarOutputStream jarOutputStream = new JarOutputStream(outputStream, targetFest);
        ) {
            while (entryEnumeration.hasMoreElements()) {
                JarEntry entry = entryEnumeration.nextElement();
                copyJarEntry(sourceJar, entry, key, jarOutputStream, handlePackages, encrypt);
            }
        }

        return targetFile;
    }

    /**
     * rename source file and copy target file to source file
     *
     * @param sourceFile source file, to be renamed
     * @param targetFile target file, replace source file
     * @param postfix    add to source file to rename
     * @return source file
     */
    protected File renameAndSwap(File sourceFile, File targetFile, String postfix) throws IOException {
        //rename source file by copying
        File renameFile = new File(sourceFile.getAbsolutePath().concat(".").concat(postfix));
        try (InputStream inputStream = new FileInputStream(sourceFile);
             OutputStream outputStream = new FileOutputStream(renameFile)) {
            IOUtils.copy(inputStream, outputStream);
        }

        //copy target file to source file
        try (InputStream inputStream = new FileInputStream(targetFile);
             OutputStream outputStream = new FileOutputStream(sourceFile)) {
            IOUtils.copy(inputStream, outputStream);
        }

        return sourceFile;
    }

    /**
     * create a new JarEntry and add it to JarOutputStream
     *
     * @param sourceFile      source jar file
     * @param entry           entry to be copied
     * @param key             encrypt key
     * @param jarOutputStream target output stream
     * @throws IOException
     */
    protected void copyJarEntry(JarFile sourceFile, JarEntry entry, String key, JarOutputStream jarOutputStream, List<String> handlePackages, boolean encrypt) throws IOException {
        String name = entry.getName();
        if (StringUtils.startsWith(name, "META-INF")) {
            return;
        }

        if (StringUtils.endsWith(name, "/")) {
            //entry is directory
            JarEntry targetEntry = new JarEntry(new ZipEntry(entry));
            addEntry(null, jarOutputStream, targetEntry);
            return;
        }

        boolean handle = false;

        if (StringUtils.endsWith(name, ".class")) {
            //entry is class file

            if (handlePackages != null && handlePackages.size() > 0) {
                String className = name.replace("/", ".");
                if (className.startsWith("BOOT-INF.classes.")) {
                    className = className.replace("BOOT-INF.classes.", "");
                }
                for (String pack : handlePackages) {
                    if (className.startsWith(pack) || className.matches(pack)) {
                        handle = true;
                        break;
                    }
                }
            } else {
                handle = true;
            }
        }

        if (handle) {
            File targetFile = null;
            if (StringUtils.contains(name, "/")) {
                String parentDir = name.substring(0, name.lastIndexOf("/"));
                String shortName = name.substring(name.lastIndexOf("/") + 1);
                targetFile = createTempTargetFile(getClassTempDir().getAbsolutePath().concat(File.separator).concat(parentDir), shortName);
            } else {
                targetFile = createTempTargetFile(getClassTempDir().getAbsolutePath(), name);
            }

            JarEntry targetEntry = getEncryptJarEntry(sourceFile, entry, targetFile, key, encrypt);
            addEntry(targetFile, jarOutputStream, targetEntry);
        } else {
            //entry is normal file
            JarEntry targetEntry = new JarEntry(new ZipEntry(entry));
            try (InputStream inputStream = sourceFile.getInputStream(entry)) {
                addInputStreamEntry(inputStream, jarOutputStream, targetEntry);
            }
        }
    }

    /**
     * create new JarEntry and encrypt entry stream to target file
     *
     * @param sourceJar source jar file
     * @param entry     class entry
     * @param key       encrypt key
     * @return new jar entry
     */
    protected JarEntry getEncryptJarEntry(JarFile sourceJar, JarEntry entry, File targetFile, String key, boolean encrypt) throws IOException {
        JarEntry targetEntry = new JarEntry(entry.getName());
        targetEntry.setTime(entry.getTime());
        targetEntry.setComment(entry.getComment());
        targetEntry.setMethod(entry.getMethod());

        try (InputStream inputStream = sourceJar.getInputStream(entry);
             OutputStream outputStream = new FileOutputStream(targetFile)) {
            byte[] bytes = IOUtils.readFully(inputStream, inputStream.available());
            byte[] targets = null;
            if (encrypt) {
                targets = AES.encrypt(bytes, AES.pwdHandler(key));
            } else {
                targets = AES.decrypt(bytes, AES.pwdHandler(key));
            }
            IOUtils.write(targets, outputStream);
        }

        return targetEntry;
    }

    /**
     * add new entry to target JarOutputStream
     *
     * @param targetFile      if not null, then copy to the target output stream
     * @param jarOutputStream target output stream
     * @param targetEntry     jar entry to be added
     */
    protected void addEntry(File targetFile, JarOutputStream jarOutputStream, JarEntry targetEntry) throws IOException {
        if (targetFile != null) {
            try (InputStream inputStream = new FileInputStream(targetFile)) {
                addInputStreamEntry(inputStream, jarOutputStream, targetEntry);
            }
        } else {
            addInputStreamEntry(null, jarOutputStream, targetEntry);
        }
    }

    /**
     * add new entry to target JarOutputStream
     *
     * @param inputStream     if not null, then copy to the target output stream
     * @param jarOutputStream target output stream
     * @param targetEntry     jar entry to be added
     */
    protected void addInputStreamEntry(InputStream inputStream, JarOutputStream jarOutputStream, JarEntry targetEntry) throws IOException {
        jarOutputStream.putNextEntry(targetEntry);
        if (inputStream != null) {
            IOUtils.copy(inputStream, jarOutputStream);
        }
        jarOutputStream.closeEntry();
    }

    /**
     * create directories in guard workspace directory
     *
     * @param directory the directory should be ensured
     */
    protected void ensureDirectory(String directory) throws IOException {
        if (StringUtils.isEmpty(directory)) {
            return;
        }

        File dirFile = new File(directory);
        if (dirFile.exists()) {
            if (dirFile.isFile()) {
                dirFile.delete();
            } else {
                return;
            }
        }

        if (dirFile.mkdirs()) {
            return;
        }
    }

    /**
     * create temporary file
     *
     * @param dir  directory
     * @param name file name
     * @return new file
     */
    protected File createTempTargetFile(String dir, String name) throws IOException {
        ensureDirectory(dir);
        File file = new File(dir, name);
        if (!file.exists()) {
            file.createNewFile();
        }

        return file;
    }

}
